백준 1202번

오늘 풀어본 문제는 백준의 1202번 문제¹이다. 문제 풀이에 사용한 언어는 Python 이다.

solved.ac 기준 CLASS

CLASS 5

문제 정보

이 문제의 내용과 조건은 다음과 같다.

문제

세계적인 도둑 상덕이는 보석점을 털기로 결심했다.

상덕이가 털 보석점에는 보석이 총 $N$ 개 있다. 각 보석은 무게 $M_i$ 와 가격 $V_i$ 를 가지고 있다. 상덕이는 가방을 $K$ 개 가지고 있고, 각 가방에 담을 수 있는 최대 무게는 $C_i$ 이다. 가방에는 최대 한 개의 보석만 넣을 수 있다.

상덕이가 훔칠 수 있는 보석의 최대 가격을 구하는 프로그램을 작성하시오.

입력

첫째 줄에 $N$ 과 $K$ 가 주어진다. $(1 \le N, K \le 300,000)$

다음 $N$ 개 줄에는 각 보석의 정보 $M_i$ 와 $V_i$ 가 주어진다. $(0 \le M_i, V_i \le 1,000,000)$

다음 $K$ 개 줄에는 가방에 담을 수 있는 최대 무게 $C_i$ 가 주어진다. $(1 \le C_i \le 100,000,000)$

모든 숫자는 양의 정수이다.

출력

첫째 줄에 상덕이가 훔칠 수 있는 보석 가격의 합의 최댓값을 출력한다.

풀이과정

1번째 시도

문제를 풀기 위해 사용한 전략은, 작은 가방부터 채우는 것이었다. 작은 보석은 작은 가방과 큰 가방 모두에 들어갈 수 있지만, 큰 보석은 큰 가방에만 들어갈 수 있기 때문에 가장 보석을 많이 담기 위해서는 작은 가방부터 채워야 하는 것이다.

보석은 C++의 priority_queue 에 해당하는 heapq 를 이용하여 무게를 기준으로 정렬한 뒤, 각 가방을 작은 것 부터 순회하며 가방 안에 들어가는 보석 중 가장 큰 것을 하나씩 고르도록 한 뒤, 총합을 최종적으로 출력하는 방식으로 문제 풀이를 시도하였다.

코드는 다음과 같이 작성하였다.

import sys
from heapq import *

N, K = map(int, sys.stdin.readline().split())

jewels = list()
for _ in range(N):
    heappush(jewels, list(map(int, sys.stdin.readline().split())))

bags = list()
for _ in range(K):
    heappush(bags, int(sys.stdin.readline()))

result = 0
fitting_jewels = list()

for bag in bags:
    if not jewels:
        break

    while True:
        if bag >= jewels[0][0]:
            heappush(fitting_jewels, -heappop(jewels)[1])
        else:
            break

    if fitting_jewels:
        result -= heappop(fitting_jewels)

print(result)

실행 결과 ‘틀렸습니다’ 가 떴다.

2번째 시도

몇 가지 테스트 케이스를 이용하여 코드를 테스트 해본 결과, 의도와는 다르게 if not jewels: break 구문이 제대로 작동하지 않는다는 것을 발견하게 되었다. 그래서 무한 반복 문에서 bag >= jewels[0][0] 조건과 묶었더니 제대로 작동하였다.

코드는 다음과 같이 수정하였다.

import sys
from heapq import *

N, K = map(int, sys.stdin.readline().split())

jewels = list()
for _ in range(N):
    heappush(jewels, list(map(int, sys.stdin.readline().split())))

bags = list()
for _ in range(K):
    heappush(bags, int(sys.stdin.readline()))

result = 0
fitting_jewels = list()

for bag in bags:
    while True:
        if jewels and bag >= jewels[0][0]:
            heappush(fitting_jewels, -heappop(jewels)[1])
        else:
            break

    if fitting_jewels:
        result -= heappop(fitting_jewels)

print(result)

실행 결과 ‘틀렸습니다’ 가 떴다.

3번째 및 4번째 시도

질문 게시판을 떠돌며 어떤 테스트 케이스를 발견했고, 그 결과 내 의도와 달리 heapq 에서 보석의 무게만을 기준으로 정렬되지 않는다는 것을 발견하게 되었다.

이를 해결하기 위해 C++ 에서 처럼 비교함수를 직접 작성할 수 있는지 찾아보았고, dataclass 를 이용하여 비슷한 걸 할 수 있다는 걸 알게되었다.

코드는 다음과 같이 수정하였다.

import sys
from heapq import *
from dataclasses import dataclass, field


@dataclass(order=True)
class Jewel:
    mass: int
    value: int = field(compare=False)


def main():
    N, K = map(int, sys.stdin.readline().split())

    jewels = list()
    for _ in range(N):
        M, V = map(int, sys.stdin.readline().split())
        heappush(jewels, Jewel(M, V))

    bags = list()
    for _ in range(K):
        bags.append(int(sys.stdin.readline()))

    bags.sort()

    result = 0
    fitting_jewels = list()

    for bag in bags:
        while True:
            if jewels and bag >= jewels[0].mass:
                heappush(fitting_jewels, -heappop(jewels).value)
            else:
                break

        if fitting_jewels:
            result -= heappop(fitting_jewels)

    print(result)


if __name__ == "__main__":
    main()

3번째 제출하는 과정에서 언어를 실수로 C++ 로 선택하여 제출하여 ‘컴파일 에러’ 가 떴다.

그 다음 제출인 4번째 제출에서는 Python 으로 제대로 선택하여 제출하자, 모든 테스트 케이스를 통과하고 정답이 나오는 것을 확인할 수 있었다.

마무리

문제를 처음에 보고 정수 범위가 넘어갈까봐 걱정되어서 처음부터 Python 사용을 시도했는데, 막상 풀고나고 보니까 다른 사람들은 C++로 많이들 하는 것 같아 조금 뻘쭘했다.

그래도 덕분에 Python 을 이용한 빠른 입출력, heapq 사용법, dataclass 등 이것저것 배울 수 있는 좋은 기회가 된 것 같다.

오늘의 PS는 여기까지!

1: https://www.acmicpc.net/problem/1202

← Previous Post Next Post →